package com.shenjies88.practice.baseframework.config;

import com.shenjies88.practice.baseframework.constant.ModelName;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

/**
 * @author by shenjies88
 * @since 2025/9/13 19:45
 */
@Configuration
public class ChatModelConfig {

    @Value("${spring.ai.ollama.base-url}")
    private String baseUrl;

    @Bean
    public OllamaChatModel qwen3ChatModel() {
        OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
        return OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                                .model(ModelName.QWEN3)
                                .build())
                .build();
    }

    @Bean
    public OllamaChatModel dsr1ChatModel() {
        OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
        return OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                                .model(ModelName.DS_R1)
                                .build())
                .build();
    }

    @Primary
    @Bean
    public OllamaChatModel dsllmChatModel() {
        OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
        return OllamaChatModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                        .model(ModelName.DS_LLM)
                        .build())
                .build();
    }

    @Bean
    public OllamaEmbeddingModel dsllmEmbedModel() {
        OllamaApi ollamaApi = OllamaApi.builder().baseUrl(baseUrl).build();
        return OllamaEmbeddingModel.builder()
                .ollamaApi(ollamaApi)
                .defaultOptions(OllamaOptions.builder()
                        .model(ModelName.DS_LLM)
                        .build())
                .build();
    }
}
